#include "leafsegmentation.h"

#include <opencv2/ml/ml.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#include <opencv2/highgui/highgui.hpp>

#include <iostream>
#include <random>

using namespace cv;

ImageMask LeafSegmentation::segmentHelperKmeans(const cv::Mat& leaf_img, int num_clusters, int num_segments)
{
    ImageMask img_mask(leaf_img.size(),num_segments);
    
    // convert to HSV
    Mat img_hsv;
    cvtColor(leaf_img,img_hsv,CV_RGB2HSV);
    int nrows = img_hsv.rows;
    int ncols = img_hsv.cols;
    
    // prepare samples using S+V features
    int num_features = 2;
    Mat samples(nrows * ncols, num_features, CV_32F);
    int index = 0;
    for (int x = 0; x < nrows; x++) {
        for (int y = 0; y < ncols; y++) {
            Vec3b& hsv = img_hsv.at<Vec3b> (x, y);
            samples.at<float> (index, 0) = hsv[1]/255.0;
            samples.at<float> (index, 1) = hsv[2]/255.0;
            index++;
        }
    }
    
    Mat labels;
    Mat centers;
    kmeans(samples,num_clusters,labels,TermCriteria(TermCriteria::COUNT +TermCriteria::EPS,EM::DEFAULT_MAX_ITERS,FLT_EPSILON),3,KMEANS_PP_CENTERS,centers);
    
//    float c0=centers.at<float>(0,0);
//    float c1=centers.at<float>(1,0);
//    float c2=centers.at<float>(2,0);
//    float c3=centers.at<float>(3,0);
//    float e0=centers.at<float>(0,1);
//    float e1=centers.at<float>(1,1);
//    float e2=centers.at<float>(2,1);
//    float e4=centers.at<float>(3,1);
    int leaf_label=0; float max_center_s=centers.at<float>(0,0);
    for(int i=1;i<num_clusters;i++){
        if(centers.at<float>(i,0)>max_center_s)
        {
            leaf_label=i;
            max_center_s=centers.at<float>(i,0);
        }
    }

    // predict
    index = 0;
    for( int x = 0; x < nrows; x++ ) {
        for( int y = 0; y < ncols; y++ ) {
            img_mask.mMask.at<uchar> (x, y) = labels.at<int>(index,1)==leaf_label?1:0;
            index++;
        }
    }
    
    return img_mask;
}

ImageMask LeafSegmentation::segmentHelperEM(const cv::Mat& leaf_img, int num_clusters, int num_segments)
{
    ImageMask img_mask(leaf_img.size(),num_segments);
    
    // convert to HSV
    Mat img_hsv;
    cvtColor(leaf_img,img_hsv,CV_RGB2HSV);
    int nrows = img_hsv.rows;
    int ncols = img_hsv.cols;
    
    // prepare samples using S+V features
    int num_features = 2;
    Mat samples(nrows * ncols, num_features, CV_8UC1);
    int index = 0;
    for (int x = 0; x < nrows; x++) {
        for (int y = 0; y < ncols; y++) {
            Vec3b& hsv = img_hsv.at<Vec3b> (x, y);
            samples.at<uchar> (index, 0) = hsv[1];
            samples.at<uchar> (index, 1) = hsv[2];
            index++;
        }
    }
    
    // initialize EM
    EM em(num_segments, EM::COV_MAT_DIAGONAL, TermCriteria(TermCriteria::COUNT +
                                                           TermCriteria::EPS,
                                                           EM::DEFAULT_MAX_ITERS,
                                                           FLT_EPSILON));
    // train
    em.train(samples);
    
    // predict
    index = 0;
    for( int x = 0; x < nrows; x++ ) {
        for( int y = 0; y < ncols; y++ ) {
            img_mask.mMask.at<uchar> (x, y) = em.predict(samples.row(index))[1];
            index++;
        }
    }
    
    return img_mask;
}

ImageMask LeafSegmentation::segment(const Mat& leaf_img, int num_clusters, int num_segments)
{
    //return segmentHelperEM(leaf_img,num_clusters,num_segments);
    return segmentHelperKmeans(leaf_img,num_clusters,num_segments);
}

int LeafSegmentation::numPointsOnBorder(const std::vector<Point> contour, Size img_size) {
    int cnt = 0;
    double boundary_portion=0.1;
    double boundary_width=(double)img_size.width*boundary_portion;
    double boundary_height=(double)img_size.height*boundary_portion;
    for (unsigned int i = 0; i < contour.size(); i++) {
        Point p = contour[i];
        if (p.x < boundary_width || p.x > img_size.width - boundary_width || p.y < boundary_height || p.y >  img_size.height  -  boundary_height)
            cnt++;
    }
    return cnt;
}


ImageMask LeafSegmentation::removeFalsePositive(const ImageMask& segmentation_mask, double threshold, int structure_size)
{
    // apply dilation first
    int dilation_size = structure_size;
    int dilation_type = MORPH_ELLIPSE;
    ImageMask dilated_mask(segmentation_mask.mMask.size(),segmentation_mask.mNumberOfLayers);
    Mat dilation_element = getStructuringElement(dilation_type,
                                                 Size(2*dilation_size + 1, 2*dilation_size+1),
                                                 Point(dilation_size, dilation_size));
    dilate(segmentation_mask.mMask, dilated_mask.mMask, dilation_element);
    
    // find connected components contour on the dilated image
    dilated_mask.connectedComponents();
    dilated_mask.componentBoundary();
    
    // exclude connected components that have large portion on border
    // and find component with largest area
    double max_area = -1.0, max_index = -1;
    std::vector<std::vector<Point> > contours = dilated_mask.mComponentBoundary;
    for (int cc = 1; cc <= (int)dilated_mask.mNumComponents; cc++) {
        std::vector<Point> contour = contours[cc];
        int num_pts_on_border = numPointsOnBorder(contour, dilated_mask.mMask.size());
        double ratio_on_border = (double)num_pts_on_border / (double)contour.size();
        if (ratio_on_border <= threshold) {
            double area = sum(dilated_mask.mComponentLabels == cc).val[0];
            if (area > max_area) { max_area = area; max_index = cc; }}
    }
    
    // take largest connected component
    ImageMask mask(segmentation_mask.mMask.size(),segmentation_mask.mNumberOfLayers);
    mask.mMask.setTo(0);
    if(max_index <= 0) return mask;
    segmentation_mask.mMask.copyTo(mask.mMask, dilated_mask.mComponentLabels == max_index);
    return mask;
}

ImageMask LeafSegmentation::removeStem(ImageMask& segmentation_mask) 
{
    // find input leaf bounding box -- use leaf size to determine erosion/dilation structure size
    segmentation_mask.connectedComponents();
    segmentation_mask.componentBoundary();
    int largest_cc = segmentation_mask.largestComponent();
    if (largest_cc <=0 || segmentation_mask.mComponentBoundary[largest_cc].size()==0) return segmentation_mask;
    RotatedRect min_box = minAreaRect(segmentation_mask.mComponentBoundary[largest_cc]);
    float leaf_size = min(min_box.size.height, min_box.size.width);
    int structure_size = (int)(leaf_size / 10);
    
    // apply top-hat
    int erosion_size = structure_size;
    int dilation_size = structure_size;
    int erosion_type = MORPH_ELLIPSE;
    int dilation_type = MORPH_ELLIPSE;
    Mat dst;
    Mat erosion_element = getStructuringElement(erosion_type,
                                                Size(2*erosion_size + 1, 2*erosion_size+1),
                                                Point(erosion_size, erosion_size));
    Mat dilation_element = getStructuringElement(dilation_type,
                                                 Size(2*dilation_size + 1, 2*dilation_size+1),
                                                 Point(dilation_size, dilation_size));
    erode(segmentation_mask.mMask, dst, erosion_element);
    dilate(dst, dst, dilation_element);
    ImageMask top_hat_mask(segmentation_mask.mMask.size(),segmentation_mask.mNumberOfLayers);
    top_hat_mask.mMask.setTo(0);
    segmentation_mask.mMask.copyTo(top_hat_mask.mMask, dst==0);
    
    // find contours of top-hat
    top_hat_mask.connectedComponents();
    top_hat_mask.componentBoundary();
    std::vector<std::vector<Point> >& contours = top_hat_mask.mComponentBoundary;
    
    // examine each connected component -- determine stem contour
    float max_ratio = 0.0;
    int stem_index = -1;
    for (int cc = 1; cc < top_hat_mask.mNumComponents; cc++) {
        std::vector<Point>& contour = contours[cc];
        if(contour.size() < 5) continue;
        // compare ratios of min area bounding box
        RotatedRect min_rect = minAreaRect(contour);
        if (min_rect.size.area() < 10) continue;
        float h = (float)min_rect.size.height;
        float w = (float)min_rect.size.width;
        float ratio = max(h/w, w/h);
        if (ratio < 2) continue;
        if (ratio > max_ratio) {
            max_ratio = ratio;
            stem_index = cc;}
    }
    if(stem_index < 0) return segmentation_mask;

    // remove the stem
    ImageMask destemed_mask(segmentation_mask);
    destemed_mask.mMask.setTo(0, top_hat_mask.mComponentLabels == stem_index);

    // take largest connected component again
    ImageMask dilated_destemed_mask(segmentation_mask.mMask.size(),segmentation_mask.mNumberOfLayers);
    Mat dilation_element_destemed = getStructuringElement(dilation_type,
                                                          Size(2*2 + 1, 2*2+1),
                                                          Point(2, 2));
    dilate(destemed_mask.mMask, dilated_destemed_mask.mMask, dilation_element_destemed);
    dilated_destemed_mask.connectedComponents();
    int largest_destemed_cc = dilated_destemed_mask.largestComponent();;
    ImageMask output_mask(segmentation_mask.mMask.size(),segmentation_mask.mNumberOfLayers);
    output_mask.mMask.setTo(0);
    segmentation_mask.mMask.copyTo(output_mask.mMask, dilated_destemed_mask.mComponentLabels == largest_destemed_cc);
    return output_mask;
}

ImageMask LeafSegmentation::resizeMask(const ImageMask& segmentation_mask, const int target_area) {

    int area = (int)sum(segmentation_mask.mMask).val[0];
    float ratio = std::sqrt((float)target_area/area);
    Size input_size = segmentation_mask.mMask.size();
    int target_width = (int)(ratio * input_size.width);
    int target_height = (int)(ratio * input_size.height);
    Size target_size(target_width, target_height);
    ImageMask resized_mask;
    resized_mask.create(target_size, segmentation_mask.mNumberOfLayers);
    resized_mask.mMask.setTo(0);
    resize(segmentation_mask.mMask, resized_mask.mMask, target_size, 0, 0, INTER_LINEAR);
    return resized_mask;
}

void LeafSegmentation::showContour(const Size& size, std::vector<std::vector<cv::Point> > contours) {
    // draw contours and min rects
    Mat drawing = Mat::zeros(size, CV_8UC3);
    for( int i = 0; i < (int)contours.size(); i++ )
    {
        // contour - blue
        drawContours( drawing, contours, i, Scalar( 255, 0, 0 ));
        // bounding rectangle - green
        Rect bounding_box = boundingRect(contours[i]);
        Point2f rect_points1[4];
        rect_points1[0] = Point2f(bounding_box.x, bounding_box.y);
        rect_points1[1] = Point2f(bounding_box.x, bounding_box.y + bounding_box.height);
        rect_points1[2] = Point2f(bounding_box.x + bounding_box.width, bounding_box.y + bounding_box.height);
        rect_points1[3] = Point2f(bounding_box.x + bounding_box.width, bounding_box.y);
        for( int j = 0; j < 4; j++ )
            line( drawing, rect_points1[j], rect_points1[(j+1)%4], Scalar( 0, 255, 0 ));
        // min rectangle - red
        RotatedRect rotated_box = minAreaRect(contours[i]);
        Point2f rect_points2[4]; rotated_box.points( rect_points2 );
        for( int j = 0; j < 4; j++ )
            line( drawing, rect_points2[j], rect_points2[(j+1)%4], Scalar( 0, 0, 255 ));
    }
    
    // Show in a window
    namedWindow( "Contours", CV_WINDOW_AUTOSIZE );
    imshow( "Contours", drawing );
    waitKey(0);
}
